package iot.disruptor;

import com.lmax.disruptor.EventHandler;
import com.lmax.disruptor.EventTranslatorOneArg;
import com.lmax.disruptor.dsl.Disruptor;
import com.lmax.disruptor.dsl.ProducerType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.LockSupport;
import java.util.function.Consumer;

/**
 * Created by dave on 18-8-11 下午10:30.
 */
public class DisruptorQueue<T extends EventData<T>> implements EventHandler<T>, EventTranslatorOneArg<T, T>, InitializingBean {
    private Disruptor<T> disruptor;
    private Consumer<List<T>> batchSaveFunc;
    private ThreadPoolTaskExecutor executor;

    private Logger logger;
    private Class<T> entityType;
    private List<T> data = new ArrayList<>();
    private int batchSize;
    private int concurrentWrite;
    private AtomicInteger writeThread = new AtomicInteger(0); // 当前正在进行的写入任务线程数

    /**
     * @param entityType        实体类类型
     * @param batchSize         批量写入最大一次写入条数，比如100。
     * @param concurrentWrite   并发写数量，比如5。
     * @param batchSaveFunc     写入数据库的方法
     * @param executor          用于写入数据库的线程池
     */
    public DisruptorQueue(Class<T> entityType,
                          int batchSize,
                          int concurrentWrite,
                          Consumer<List<T>> batchSaveFunc,
                          ThreadPoolTaskExecutor executor) {
        this.entityType = entityType;
        this.batchSize = batchSize;
        this.concurrentWrite = concurrentWrite;
        this.batchSaveFunc = batchSaveFunc;
        this.executor = executor;
        this.logger = LoggerFactory.getLogger(DisruptorQueue.class.getSimpleName() + "-" + entityType.getSimpleName());

        disruptor = new Disruptor<>(
                () -> {
                    try {
                        return entityType.newInstance();
                    } catch (InstantiationException | IllegalAccessException e) {
                        throw new RuntimeException(e);
                    }
                },
                2 << 12,
                (Runnable r) -> {
                    Thread t = new Thread(r);
                    t.setName("DisruptorQueue-" + entityType.getSimpleName());
                    t.setDaemon(true);
                    return t;
                },
                ProducerType.MULTI,
                new TimedWaitStrategy(this, 50, 1000_000, 200));

        disruptor.handleEventsWith(this);
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        disruptor.start();
    }

    public Class<T> getEntityType() {
        return entityType;
    }

    public void stop() {
        disruptor.shutdown();
    }

    public void push(T item) {
        disruptor.publishEvent(this, item);
    }

    public void flush() {
        if (data.size() == 0) {
            return;
        }

        int i = 0;
        while (writeThread.get() >= concurrentWrite) {
            if (++i >= 100) {
                logger.warn("丢弃：" + data.size() + "条");
                data.clear(); // 100次10ms后直接清空
                return;
            }
            LockSupport.parkNanos(100_000); // 等待0.1ms
        }
        writeThread.incrementAndGet();
        final List<T> list = new ArrayList<>(data);
        data.clear();
        executor.submitListenable(() -> batchSaveFunc.accept(list)).completable().thenApplyAsync(ret -> {
            writeThread.decrementAndGet();
            return null;
        });
    }

    public long getPoolSize() {
        return disruptor.getBufferSize();
    }

    public long getQueueSize() {
        return disruptor.getBufferSize() - disruptor.getRingBuffer().remainingCapacity();
    }

    @Override
    public void onEvent(T event, long sequence, boolean endOfBatch) throws Exception {
        if (data.size() >= this.batchSize) {
            flush();
        }
        data.add(event.duplicate());
    }

    @Override
    public void translateTo(T event, long sequence, T arg) {
        arg.copyTo(event);
    }
}
